package com.origin.niuke.simulation;

import java.util.Arrays;

/**
 * 对角线遍历矩阵
 * 给定一个大小为 n*m 的矩阵，请以对角线遍历并返回遍历结果
 * 算法：模拟
 * 有六种情况(n * m)
 * 向前走（右上）：--r, ++c
 * - r = -1, c = m
 * - r = -1, c < m
 * - r < n, c = m
 * 向下走（左下）：++r, --c
 * - c = -1, r = n
 * - c = -1, r < n
 * - c < m, r = n
 *
 * @author yezh
 * @date 2023/2/14 20:37
 */
public class NC201 {

    public static void main(String[] args) {
        int[][] mat = {{1, 2, 5}, {3, 4, 6}};
        System.out.println(Arrays.toString(new NC201().diagonalOrder(mat)));
    }

    public int[] diagonalOrder(int[][] mat) {
        // write code here
        int n = mat.length, m = mat[0].length, idx = 0;
        boolean flag = true;
        int[] ans = new int[n * m];
        for (int r = 0, c = 0; idx < ans.length; ++idx) {
            ans[idx] = mat[r][c];
            if (flag) {
                --r;
                ++c;
                if (r == -1 && c < m) {
                    r = 0;
                    flag = false;
                } else if (r == -1 && c == m) {
                    r = 1;
                    --c;
                    flag = false;
                } else if (r < n && c == m) {
                    --c;
                    r += 2;
                    flag = false;
                }
            } else {
                ++r;
                --c;
                if (c == -1 && r < n) {
                    c = 0;
                    flag = true;
                } else if (c == -1 && r == n) {
                    c = 1;
                    --r;
                    flag = true;
                } else if (c < m && r == n) {
                    --r;
                    c += 2;
                    flag = true;
                }
            }
        }
        return ans;
    }

}
